package qmlt.learning.decisiontree.control;

import java.util.Collection;
import java.util.List;

import qmlt.dataset.Instance;

public class GainRatioTreeBuildController extends InfoGainTreeBuildController
{
    @Override
    protected float getSplitScore(int size, float entropy, Collection<List<Instance>> subsets)
    {
        float infoGain = super.getSplitScore(size, entropy, subsets);
        float splitInfo = getSplitInfo(size, subsets);
        return infoGain / splitInfo;
    }

    protected float getSplitInfo(int size, Collection<List<Instance>> subsets)
    {
        float splitInfo = 0;
        for (List<Instance> subset : subsets)
        {
            float ratio = (float) subset.size() / size;
            splitInfo -= ratio * Math.log(ratio) / Math.log(2);
        }
        return splitInfo;
    }
}
